"""
K线数据处理模块
"""
import pandas as pd
import numpy as np
from typing import List, Dict, Optional
from datetime import datetime, timedelta
from src.okx.models import KlineData
from src.utils.logger import get_logger

logger = get_logger(__name__)


class KlineProcessor:
    """K线数据处理器"""

    def __init__(self):
        self.cache: Dict[str, pd.DataFrame] = {}

    def add_klines(self, symbol: str, timeframe: int, klines: List[KlineData]) -> pd.DataFrame:
        """
        添加K线数据到缓存

        Args:
            symbol: 交易对
            timeframe: 时间框架
            klines: K线数据列表

        Returns:
            DataFrame格式的K线数据
        """
        key = f"{symbol}_{timeframe}"

        # 转换为DataFrame
        data = {
            'timestamp': [k.timestamp for k in klines],
            'open': [k.open for k in klines],
            'high': [k.high for k in klines],
            'low': [k.low for k in klines],
            'close': [k.close for k in klines],
            'volume': [k.volume for k in klines],
            'quote_volume': [k.quote_volume for k in klines],
        }

        df = pd.DataFrame(data)
        df['datetime'] = pd.to_datetime(df['timestamp'], unit='ms')
        df = df.sort_values('timestamp').reset_index(drop=True)

        # 如果缓存中已有数据，合并并去重
        if key in self.cache:
            existing_df = self.cache[key]
            df = pd.concat([existing_df, df], ignore_index=True)
            df = df.drop_duplicates(subset=['timestamp'], keep='last')
            df = df.sort_values('timestamp').reset_index(drop=True)

        self.cache[key] = df
        logger.debug(f"K线数据添加成功: {symbol} {timeframe}分钟, 共{len(df)}条")

        return df

    def get_klines(self, symbol: str, timeframe: int) -> Optional[pd.DataFrame]:
        """获取缓存的K线数据"""
        key = f"{symbol}_{timeframe}"
        return self.cache.get(key)

    def calculate_sma(self, df: pd.DataFrame, period: int, column: str = 'close') -> np.ndarray:
        """
        计算简单移动平均线 (SMA)

        Args:
            df: K线数据
            period: 周期
            column: 列名

        Returns:
            SMA数组
        """
        return df[column].rolling(window=period).mean().values

    def calculate_ema(self, df: pd.DataFrame, period: int, column: str = 'close') -> np.ndarray:
        """
        计算指数移动平均线 (EMA)

        Args:
            df: K线数据
            period: 周期
            column: 列名

        Returns:
            EMA数组
        """
        return df[column].ewm(span=period, adjust=False).mean().values

    def calculate_rsi(self, df: pd.DataFrame, period: int = 14, column: str = 'close') -> np.ndarray:
        """
        计算相对强弱指标 (RSI)

        Args:
            df: K线数据
            period: 周期
            column: 列名

        Returns:
            RSI数组 (0-100)
        """
        prices = df[column].values
        deltas = np.diff(prices)

        # 创建与prices长度相同的gain和loss数组
        gains = np.zeros(len(prices))
        losses = np.zeros(len(prices))

        # 填充gain和loss数组 (第0个元素为0，因为没有前一个价格)
        for i in range(1, len(prices)):
            if deltas[i-1] > 0:
                gains[i] = deltas[i-1]
            elif deltas[i-1] < 0:
                losses[i] = -deltas[i-1]

        # 初始化输出数组
        rsi = np.zeros_like(prices, dtype=float)
        rsi[:] = np.nan

        if len(prices) <= period:
            return rsi

        # 计算初始平均值
        avg_gain = gains[1:period+1].mean()
        avg_loss = losses[1:period+1].mean()

        # 在period位置设置第一个RSI值
        if avg_loss == 0:
            rsi[period] = 100 if avg_gain > 0 else 50
        else:
            rs = avg_gain / avg_loss
            rsi[period] = 100 - (100 / (1 + rs))

        # EMA计算
        for i in range(period + 1, len(prices)):
            avg_gain = (avg_gain * (period - 1) + gains[i]) / period
            avg_loss = (avg_loss * (period - 1) + losses[i]) / period

            if avg_loss == 0:
                rsi[i] = 100 if avg_gain > 0 else 50
            else:
                rs = avg_gain / avg_loss
                rsi[i] = 100 - (100 / (1 + rs))

        return rsi

    def calculate_macd(self, df: pd.DataFrame, fast: int = 12, slow: int = 26,
                      signal: int = 9, column: str = 'close') -> tuple:
        """
        计算MACD指标

        Args:
            df: K线数据
            fast: 快线周期
            slow: 慢线周期
            signal: 信号线周期
            column: 列名

        Returns:
            (MACD线, 信号线, 直方图) 元组
        """
        ema_fast = self.calculate_ema(df, fast, column)
        ema_slow = self.calculate_ema(df, slow, column)

        macd_line = ema_fast - ema_slow
        signal_line = pd.Series(macd_line).ewm(span=signal, adjust=False).mean().values
        histogram = macd_line - signal_line

        return macd_line, signal_line, histogram

    def calculate_bollinger_bands(self, df: pd.DataFrame, period: int = 20,
                                 std_dev: float = 2.0, column: str = 'close') -> tuple:
        """
        计算布林带

        Args:
            df: K线数据
            period: 周期
            std_dev: 标准差倍数
            column: 列名

        Returns:
            (中线, 上线, 下线) 元组
        """
        middle = df[column].rolling(window=period).mean().values
        std = df[column].rolling(window=period).std().values
        upper = middle + (std * std_dev)
        lower = middle - (std * std_dev)

        return middle, upper, lower

    def calculate_atr(self, df: pd.DataFrame, period: int = 14) -> np.ndarray:
        """
        计算平均真实波幅 (ATR)

        Args:
            df: K线数据
            period: 周期

        Returns:
            ATR数组
        """
        high = df['high'].values
        low = df['low'].values
        close = df['close'].values

        # 计算真实波幅 (TR)
        tr1 = high - low
        tr2 = np.abs(high - np.roll(close, 1))
        tr3 = np.abs(low - np.roll(close, 1))

        tr = np.maximum(tr1, np.maximum(tr2, tr3))

        # 计算ATR (简单移动平均)
        atr = pd.Series(tr).rolling(window=period).mean().values

        return atr

    def add_technical_indicators(self, df: pd.DataFrame, config: Dict) -> pd.DataFrame:
        """
        为DataFrame添加技术指标

        Args:
            df: K线数据
            config: 策略配置

        Returns:
            添加指标后的DataFrame
        """
        if len(df) < 30:
            logger.warning(f"数据不足, 无法计算技术指标")
            return df

        # RSI
        df['rsi'] = self.calculate_rsi(df, config.get('rsi_period', 14))

        # MACD
        macd, signal, histogram = self.calculate_macd(
            df,
            config.get('macd_fast', 12),
            config.get('macd_slow', 26),
            config.get('macd_signal', 9)
        )
        df['macd'] = macd
        df['macd_signal'] = signal
        df['macd_histogram'] = histogram

        # 布林带
        middle, upper, lower = self.calculate_bollinger_bands(
            df,
            period=20,
            std_dev=config.get('bollinger_std', 2.0)
        )
        df['bb_middle'] = middle
        df['bb_upper'] = upper
        df['bb_lower'] = lower

        # ATR (波幅)
        df['atr'] = self.calculate_atr(df)

        # 价格变化率
        df['pct_change'] = df['close'].pct_change() * 100

        return df

    def get_recent_bars(self, symbol: str, timeframe: int, n: int = 20) -> Optional[pd.DataFrame]:
        """获取最近N根K线"""
        df = self.get_klines(symbol, timeframe)
        if df is None or len(df) == 0:
            return None

        return df.tail(n).copy()
